# 本code可参考：https://blog.csdn.net/m0_37579232/article/details/119879646
# paper: http://ydwen.github.io/papers/WenECCV16.pdf
#  code:  https://github.com/pangyupo/mxnet_center_loss
#  pytorch code: https://blog.csdn.net/sinat_37787331/article/details/80296964

import torch
import torch.nn as nn
from torch.autograd.function import Function
from ..builder import LOSSES

@LOSSES.register_module()
class CenterLoss(nn.Module):
    def __init__(self, num_classes, feat_dim=2,  loss_weight=1.0, size_average=True):
        super(CenterLoss, self).__init__()
        self.centers = nn.Parameter(torch.randn(num_classes, feat_dim))
        self.centerlossfunc = CenterlossFunc.apply
        self.feat_dim = feat_dim
        self.size_average = size_average
        self.loss_weight = loss_weight

    def forward(self,
                label,
                feat,
                weight=None,
                avg_factor=None,  # 针对报错的解决方法
                reduction_override=None,
                ):
        batch_size = feat.size(0)
        feat = feat.view(batch_size, -1)
        # To check the dim of centers and features
        if feat.size(1) != self.feat_dim:
            raise ValueError("Center's dim: {0} should be equal to input feature's \
                            dim: {1}".format(self.feat_dim, feat.size(1)))
        batch_size_tensor = feat.new_empty(1).fill_(batch_size if self.size_average else 1)
        loss = self.loss_weight * self.centerlossfunc(feat, label, self.centers, batch_size_tensor)
        return loss


class CenterlossFunc(Function):
    @staticmethod
    def forward(ctx, feature, label, centers, batch_size):
        ctx.save_for_backward(feature, label, centers, batch_size)
        centers_batch = centers.index_select(0, label.long())
        return (feature - centers_batch).pow(2).sum() / 2.0 / batch_size

    @staticmethod
    def backward(ctx, grad_output):
        feature, label, centers, batch_size = ctx.saved_tensors
        centers_batch = centers.index_select(0, label.long())
        diff = centers_batch - feature
        # init every iteration
        counts = centers.new_ones(centers.size(0))
        ones = centers.new_ones(label.size(0))
        grad_centers = centers.new_zeros(centers.size())

        counts = counts.scatter_add_(0, label.long(), ones)
        grad_centers.scatter_add_(0, label.unsqueeze(1).expand(feature.size()).long(), diff)
        grad_centers = grad_centers/counts.view(-1, 1)
        return - grad_output * diff / batch_size, None, grad_centers / batch_size, None


# def main(test_cuda=False):
#     print('-'*80)
#     device = torch.device("cuda" if test_cuda else "cpu")
#     ct = CenterLoss(10,2,size_average=True).to(device)
#     y = torch.Tensor([0,0,2,1]).to(device)
#     feat = torch.zeros(4,2).to(device).requires_grad_()
#     print (list(ct.parameters()))
#     print (ct.centers.grad)
#     out = ct(y,feat)
#     print(out.item())
#     out.backward()
#     print(ct.centers.grad)
#     print(feat.grad)

# if __name__ == '__main__':
#     torch.manual_seed(999)
#     main(test_cuda=False)
#     if torch.cuda.is_available():
#         main(test_cuda=True)

# 报错： forward() got an unexpected keyword argument 'avg_factor'
